#!usr/bin/env python
# -*- coding: utf-8 -*-

"""重分类
根据给定的 csv 格式重分类表的范围, 将 Tiff 文件重分类。

Example:
    $ python reclassify.py -o out.tif source.tif classify.csv
"""

import sys
import csv
import numpy as np
from osgeo import gdal, ogr, osr

def raster2array(raster_fn):
    raster = gdal.Open(raster_fn)
    band = raster.GetRasterBand(1)
    array = band.ReadAsArray()
    return array

def no_data_value(raster_fn):
    raster = gdal.Open(raster_fn)
    band = raster.GetRasterBand(1)
    no_data = band.GetNoDataValue()
    return no_data

def array2raster(raster_fn, array, new_raster_fn):
    raster = gdal.Open(raster_fn)
    geotransform = raster.GetGeoTransform()
    originX, pixelWidth, _, originY, _, pixelHeight = geotransform
    cols = array.shape[1]
    rows = array.shape[0]

    driver = gdal.GetDriverByName('GTiff')
    outRaster = driver.Create(new_raster_fn, cols, rows, 1, gdal.GDT_Int32)
    outRaster.SetGeoTransform((originX, pixelWidth, 0, originY, 0, pixelHeight))
    outband = outRaster.GetRasterBand(1)
    outband.SetNoDataValue(no_data_value(raster_fn))
    outband.WriteArray(array)
    outRasterSRS = osr.SpatialReference()
    outRasterSRS.ImportFromWkt(raster.GetProjectionRef())
    outRaster.SetProjection(outRasterSRS.ExportToWkt())
    outband.FlushCache()

def main(raster_fn, csv_file, new_raster_fn):
    classify_table = []
    with open(csv_file) as csvfile:
        spamreader = csv.reader(csvfile, delimiter=',')
        for row in spamreader:
            classify_table.append([float(i) for i in row])

    array = raster2array(raster_fn)
    for row in classify_table:
        min_val, max_val, new_val = row
        array[np.where((array>min_val) & (array<=max_val))] = new_val

    array2raster(raster_fn, array, new_raster_fn)

if __name__ == "__main__":
    argv = sys.argv

    raster_fn = None
    csv_fn = None
    output_fn = None

    i = 1
    while i < len(argv):
        arg = argv[i]

        if arg == '-o' or arg == '-out':
            i = i + 1
            output_fn = argv[i]

        elif raster_fn is None:
            raster_fn = argv[i]

        elif csv_fn is None:
            csv_fn = argv[i]

        else:
            print("python script.py -o out.tif source.tif classify.csv")
            sys.exit(1)

        i = i + 1

    if output_fn is None:
        output_fn = 'OUT.tif'

    main(raster_fn, csv_fn, output_fn)